########################################################################################
######################          Import packages      ###################################
########################################################################################
from flask import Blueprint, render_template, redirect, url_for, request, flash, jsonify
from werkzeug.security import generate_password_hash, check_password_hash
from models import User
from flask_login import login_user, logout_user, login_required
from __init__ import db

import base64
from Crypto import Random
from Crypto.Cipher import PKCS1_v1_5 as Cipher_pkcs1_v1_5
from Crypto.PublicKey import RSA
import random
import hashlib


auth = Blueprint('auth', __name__) # 为身份鉴别系统单独创建一个蓝图

# 生成RSA密钥
# 产生的密钥已经存在文件目录中，需要定期更换
# private.pem public.pem
key = RSA.generate(2048)
private_key = key.export_key()
public_key = key.publickey().export_key()

# 进行质询的路由
@auth.route('/do_challenge', methods=['POST'])
def do_challenge():
    # 读取公钥，准备发送
    with open('public.pem', "r") as f: 
        text=f.read().splitlines() # 获取文件全部数据，不要回车, 返回结果是一个列表
    public_key = ""                # 这里是为了匹配前端js加密时读取密钥不到回车，所以在后端直接进行过滤回车的处理
    for t in text:
        public_key = public_key + t
    print(public_key)

    # 产生质询数 challenge
    global challenge
    challenge = str(random.randint(10000, 99999))
    # print("产生的质询数是：%s", challenge)

    return jsonify(public_key=public_key, challenge=challenge)

# 进行验证的路由
@auth.route('/login', methods=['GET', 'POST'])
def login(): 
    if request.method=='GET': # GET请求返回登录页面
        return render_template('login.html')
    else: # POST请求说明前端发来验证
        
        # 读取私钥
        private_key = ''
        with open('private.pem', 'r') as f:
            private_key = f.read()

        # 读取公钥
        public_key = ''
        with open('public.pem', 'r') as f:
            public_key = f.read()
        publicKey = RSA.import_key(public_key)

        # 生成server_id
        with open('public.pem', "r") as f: 
            text=f.read().splitlines() # 验证者标识符由验证者的公钥进行MD5哈希计算得到
        server_id = ""
        for t in text:
            server_id = server_id + t
        hl = hashlib.md5()
        hl.update(server_id.encode(encoding='utf-8'))
        server_id = hl.hexdigest()
        # print("生成的%s", hl.hexdigest())

        # 对token进行解密
        encrypted_data = request.form.get('token')
        print(encrypted_data)
        privateKey = RSA.import_key(private_key)
        cipher = Cipher_pkcs1_v1_5.new(privateKey)
        random_generator = Random.new().read
        data = cipher.decrypt(base64.b64decode(encrypted_data), None)
        token = data.decode('utf-8')

        # 分割token
        result = token.split("|")
        email = result[0]
        password = result[1]
        client_challenge = result[2]
        client_server_id = result[3]
        # print("发来质询数是%s", client_challenge)
        # print(email)
        # print(password)
        # print(server_id)
        
        # 查看用户是否存在
        remember = True if request.form.get('remember') else False
        user = User.query.filter_by(email=email).first()

        # 用户不存在
        if not user:
            flash('Please sign up before!')
            return redirect(url_for('auth.signup'))
        
        # 密码不对
        elif not check_password_hash(user.password, password):
            flash('Please check your password and try again.')
            return redirect(url_for('auth.login')) 
        
        # 质询数不对
        elif(str(client_challenge) != str(challenge)):
            flash('I think you are attacking!')
            return redirect(url_for('auth.login')) 
        
        # server_id不对
        elif(str(client_server_id) != str(server_id)):
            flash('Your RSA public key(server id) is wrong!')
            return redirect(url_for('auth.login')) 
        
        # 如果验证都通过了，成功登录
        login_user(user, remember=remember)
        # challenge = str(random.randint(10000, 99999))
        return redirect(url_for('main.profile'))

@auth.route('/signup', methods=['GET', 'POST'])
def signup(): 
    if request.method=='GET': # GET请求返回注册页面
        return render_template('signup.html')
    else: # POST请求说明前端发来注册申请

        # 这里后来决定取消对非password的加密，对用户名没有必要加密
        email = request.form.get('encryptedEmail')
        name = request.form.get('encryptedName')
        password = request.form.get('encryptedPassword')    

        # 查看是否已经注册过了
        user = User.query.filter_by(email=email).first() 
        if user: 
            flash('Email address already exists')
            return redirect(url_for('auth.signup'))
        
        # 建立一个新用户，密码采用哈希值存储，不能明文存储
        new_user = User(email=email, name=name, password=generate_password_hash(password)) #
        # 添加到数据库
        db.session.add(new_user)
        db.session.commit()
        return redirect(url_for('auth.login'))

@auth.route('/logout') # 退出登录的路由
@login_required # 设定登录权限
def logout(): 
    logout_user()
    return redirect(url_for('main.index'))